/*
 *   Copyright (c) 2014, Andrew Romanenko <melanhit@gmail.com>
 *   All rights reserved.
 *
 *   Redistribution and use in source and binary forms, with or without
 *   modification, are permitted provided that the following conditions are met:
 *
 *   1. Redistributions of source code must retain the above copyright notice, this
 *      list of conditions and the following disclaimer.
 *   2. Redistributions in binary form must reproduce the above copyright notice,
 *      this list of conditions and the following disclaimer in the documentation
 *      and/or other materials provided with the distribution.
 *   3. Neither the name of the project nor the names of its contributors
 *      may be used to endorse or promote products derived from this software
 *      without specific prior written permission.
 *
 *   THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND
 *   ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 *   WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *   DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 *   ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 *   (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *   LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 *   ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 *   (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 *   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include "akmos.h"
#include "cipher.h"

#include "algo/anubis.h"
#include "algo/cast6.h"
#include "algo/rc6.h"
#include "algo/serpent.h"
#include "algo/twofish.h"
#include "algo/threefish.h"
#include "algo/camellia.h"

#include "mode/ecb.h"
#include "mode/cbc.h"
#include "mode/ofb.h"
#include "mode/ctr.h"

static int cipher_init_algo(akmos_cipher_t *ctx, akmos_algo_t algo)
{
    int err;

    err = AKMOS_ERR_SUCCESS;

    switch(algo) {
        case AKMOS_ALGO_ANUBIS:
            ctx->algo  = &akmos_algo_anubis;
            ctx->actx  = malloc(sizeof(akmos_anubis_t));
            break;

        case AKMOS_ALGO_CAST6:
            ctx->algo  = &akmos_algo_cast6;
            ctx->actx  = malloc(sizeof(akmos_cast6_t));
            break;

        case AKMOS_ALGO_RC6:
            ctx->algo  = &akmos_algo_rc6;
            ctx->actx  = malloc(sizeof(akmos_rc6_t));
            break;

        case AKMOS_ALGO_SERPENT:
            ctx->algo  = &akmos_algo_serpent;
            ctx->actx  = malloc(sizeof(akmos_serpent_t));
            break;

        case AKMOS_ALGO_TWOFISH:
            ctx->algo  = &akmos_algo_twofish;
            ctx->actx  = malloc(sizeof(akmos_twofish_t));
            break;

        case AKMOS_ALGO_THREEFISH_256:
            ctx->algo  = &akmos_algo_threefish_256;
            ctx->actx  = malloc(sizeof(akmos_threefish_256_t));
            break;

        case AKMOS_ALGO_THREEFISH_512:
            ctx->algo  = &akmos_algo_threefish_512;
            ctx->actx  = malloc(sizeof(akmos_threefish_512_t));
            break;

        case AKMOS_ALGO_THREEFISH_1024:
            ctx->algo  = &akmos_algo_threefish_1024;
            ctx->actx  = malloc(sizeof(akmos_threefish_1024_t));
            break;

        case AKMOS_ALGO_CAMELLIA:
            ctx->algo  = &akmos_algo_camellia;
            ctx->actx  = malloc(sizeof(akmos_camellia_t));
            break;

        default:
            err = AKMOS_ERR_ALGOID;
            goto out;
    }

    if(!ctx->actx) {
         err = AKMOS_ERR_ENOMEM;
         goto out;
    }

    return AKMOS_ERR_SUCCESS;

/* process errors */
out:
    if(ctx->actx)
        free(ctx->actx);

    return err;
}

static int cipher_init_mode(akmos_cipher_t *ctx, akmos_mode_t mode)
{
    int err;

    err = AKMOS_ERR_SUCCESS;

    switch(mode) {
        case AKMOS_MODE_ECB:
            ctx->encrypt = (void *)&akmos_ecb_encrypt;
            ctx->decrypt = (void *)&akmos_ecb_decrypt;
            break;

        case AKMOS_MODE_CBC:
            ctx->init    = (void *)&akmos_cbc_init;
            ctx->setiv   = (void *)&akmos_cbc_setiv;
            ctx->encrypt = (void *)&akmos_cbc_encrypt;
            ctx->decrypt = (void *)&akmos_cbc_decrypt;
            ctx->free    = (void *)&akmos_cbc_free;
            ctx->mctx    = malloc(sizeof(akmos_cbc_t));
            break;

        case AKMOS_MODE_OFB:
            ctx->init    = (void *)&akmos_ofb_init;
            ctx->setiv   = (void *)&akmos_ofb_setiv;
            ctx->encrypt = (void *)&akmos_ofb_encrypt;
            ctx->decrypt = (void *)&akmos_ofb_encrypt;
            ctx->free    = (void *)&akmos_ofb_free;
            ctx->mctx    = malloc(sizeof(akmos_ofb_t));
            break;

        case AKMOS_MODE_CTR:
            ctx->init    = (void *)&akmos_ctr_init;
            ctx->setiv   = (void *)&akmos_ctr_setiv;
            ctx->encrypt = (void *)&akmos_ctr_encrypt;
            ctx->decrypt = (void *)&akmos_ctr_encrypt;
            ctx->free    = (void *)&akmos_ctr_free;
            ctx->mctx    = malloc(sizeof(akmos_ctr_t));
            break;

        default:
            err = AKMOS_ERR_MODEID;
            goto out;
    }

    ctx->id = mode;

    if(ctx->init) {
        if(!ctx->mctx) {
            err = AKMOS_ERR_ENOMEM;
            goto out;
        }

        err = ctx->init(ctx);
        if(err)
            goto out;
    }

    return AKMOS_ERR_SUCCESS;

/* process errors */
out:
    if(ctx->mctx)
        free(ctx->mctx);

    return err;
}

int akmos_cipher_init(akmos_cipher_t **ctx, akmos_algo_t algo, akmos_mode_t mode)
{
    akmos_cipher_t *ptr;
    int err;

    err = AKMOS_ERR_SUCCESS;

    ptr = *ctx = malloc(sizeof(akmos_cipher_t));
    if(!ptr)
        return AKMOS_ERR_ENOMEM;

    memset(ptr, 0, sizeof(akmos_cipher_t));

    err = cipher_init_algo(ptr, algo);
    if(err)
        goto out;

    err = cipher_init_mode(ptr, mode);
    if(err)
        goto out;

    return AKMOS_ERR_SUCCESS;

/* process errors */
out:
    if(ptr)
        free(ptr);

    return err;
}

int akmos_cipher_setkey(akmos_cipher_t *ctx, const uint8_t *key, size_t len)
{
    if(len < ctx->algo->keymin || len > ctx->algo->keymax)
        return AKMOS_ERR_KEYLEN;

    if((len % ctx->algo->keystep) != 0)
        return AKMOS_ERR_KEYLEN;

    ctx->algo->setkey(ctx->actx, key, len);

    return AKMOS_ERR_SUCCESS;
}

void akmos_cipher_setiv(akmos_cipher_t *ctx, const uint8_t *iv)
{
    switch(ctx->id) {
        case AKMOS_MODE_ECB:
            break;

        case AKMOS_MODE_CBC:
        case AKMOS_MODE_OFB:
        case AKMOS_MODE_CTR:
            ctx->setiv(ctx, iv);
            break;

        default:
            break;
    }
}

void akmos_cipher_encrypt(akmos_cipher_t *ctx, const uint8_t *in_blk, size_t in_len, uint8_t *out_blk)
{
    ctx->encrypt(ctx, in_blk, in_len, out_blk);
}

void akmos_cipher_decrypt(akmos_cipher_t *ctx, const uint8_t *in_blk, size_t in_len, uint8_t *out_blk)
{
    ctx->decrypt(ctx, in_blk, in_len, out_blk);
}

void akmos_cipher_free(akmos_cipher_t *ctx)
{
    int t_len;

    if(!ctx)
        return;

    switch(ctx->algo->id) {
        case AKMOS_ALGO_ANUBIS:
            t_len = sizeof(akmos_anubis_t);
            break;

        case AKMOS_ALGO_CAST6:
            t_len = sizeof(akmos_cast6_t);
            break;

        case AKMOS_ALGO_RC6:
            t_len = sizeof(akmos_rc6_t);
            break;

        case AKMOS_ALGO_SERPENT:
            t_len = sizeof(akmos_serpent_t);
            break;

        case AKMOS_ALGO_TWOFISH:
            t_len = sizeof(akmos_twofish_t);
            break;

        case AKMOS_ALGO_THREEFISH_256:
            t_len = sizeof(akmos_threefish_256_t);
            break;

        case AKMOS_ALGO_THREEFISH_512:
            t_len = sizeof(akmos_threefish_512_t);
            break;

        case AKMOS_ALGO_THREEFISH_1024:
            t_len = sizeof(akmos_threefish_1024_t);
            break;

        case AKMOS_ALGO_CAMELLIA:
            t_len = sizeof(akmos_camellia_t);
            break;

        default:
            return;
    }

    memset(ctx->actx, 0, t_len);
    free(ctx->actx);

    switch(ctx->id) {
        case AKMOS_MODE_ECB:
            break;

        case AKMOS_MODE_CBC:
        case AKMOS_MODE_OFB:
        case AKMOS_MODE_CTR:
            ctx->free(ctx);
            break;

        default:
            return;
    }

    memset(ctx, 0, sizeof(akmos_cipher_t));
    free(ctx);
}

int akmos_cipher_ex(int enc, akmos_algo_t algo, akmos_mode_t mode, const uint8_t *key, size_t keylen,
                    const uint8_t *iv, const uint8_t *in_blk, size_t in_len, uint8_t *out_blk)
{
    akmos_cipher_t *ctx;
    int err;

    err = akmos_cipher_init(&ctx, algo, mode);
    if(err)
        goto out;

    err = akmos_cipher_setkey(ctx, key, keylen);
    if(err)
        goto out;

    switch(mode) {
        case AKMOS_MODE_CBC:
        case AKMOS_MODE_OFB:
        case AKMOS_MODE_CTR:
            akmos_cipher_setiv(ctx, iv);
            break;

        default:
            break;
    }

    if(enc)
        akmos_cipher_encrypt(ctx, in_blk, in_len, out_blk);
    else
        akmos_cipher_decrypt(ctx, in_blk, in_len, out_blk);

out:
    if(ctx)
        akmos_cipher_free(ctx);

    return err;
}
